package com.example.batch.core.transaction;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.support.DefaultTransactionStatus;
import javax.sql.DataSource;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 分布式事务管理器
 * 支持多数据源的分布式事务处理,包括XA事务和最大努力事务
 */
@Slf4j
@Component
public class DistributedTransactionManager implements PlatformTransactionManager {

    // 存储每个数据源对应的事务管理器
    private final Map<String, DataSourceTransactionManager> transactionManagers;
    
    // XA数据源包装器,用于处理分布式事务
    private final XADataSourceWrapper xaWrapper;
    
    // 事务同步管理器
    private final TransactionSynchronizationManager synchronizationManager;
    
    // 事务超时时间(秒)
    private final int defaultTimeout;

    /**
     * 构造函数
     * @param dynamicDataSource 动态数据源
     * @param xaWrapper XA数据源包装器
     * @param properties 事务配置属性
     */
    public DistributedTransactionManager(
            DynamicRoutingDataSource dynamicDataSource,
            XADataSourceWrapper xaWrapper,
            TransactionProperties properties) {
        this.xaWrapper = xaWrapper;
        this.defaultTimeout = properties.getTimeout();
        this.synchronizationManager = new TransactionSynchronizationManager();
        this.transactionManagers = initTransactionManagers(dynamicDataSource);
    }

    /**
     * 获取事务,根据定义创建或加入已存在的事务
     * @param definition 事务定义
     * @return 事务状态
     */
    @Override
    public TransactionStatus getTransaction(TransactionDefinition definition) {
        // 获取当前数据源名称
        String dataSourceName = DynamicRoutingDataSource.getDataSource();
        if (dataSourceName == null) {
            throw new IllegalStateException("No active data source found");
        }
        
        // 获取对应的事务管理器
        PlatformTransactionManager targetTM = transactionManagers.get(dataSourceName);
        if (targetTM == null) {
            throw new IllegalStateException("No transaction manager found for: " + dataSourceName);
        }
        
        try {
            // 初始化事务同步
            synchronizationManager.initSynchronization();
            
            // 创建事务
            TransactionStatus status = targetTM.getTransaction(definition);
            
            // 记录事务开始
            log.debug("Started transaction for data source: {}", dataSourceName);
            
            return status;
        } catch (Exception e) {
            // 清理同步状态
            synchronizationManager.clearSynchronization();
            throw new TransactionSystemException("Could not create transaction", e);
        }
    }

    /**
     * 提交事务
     * @param status 事务状态
     */
    @Override
    public void commit(TransactionStatus status) {
        try {
            String dataSourceName = DynamicRoutingDataSource.getDataSource();
            PlatformTransactionManager targetTM = transactionManagers.get(dataSourceName);
            
            if (targetTM == null) {
                throw new IllegalStateException("No transaction manager found for commit");
            }
            
            targetTM.commit(status);
            log.debug("Committed transaction for data source: {}", dataSourceName);
        } finally {
            cleanup(status);
        }
    }

    /**
     * 回滚事务
     * @param status 事务状态
     */
    @Override
    public void rollback(TransactionStatus status) {
        try {
            String dataSourceName = DynamicRoutingDataSource.getDataSource();
            PlatformTransactionManager targetTM = transactionManagers.get(dataSourceName);
            
            if (targetTM == null) {
                throw new IllegalStateException("No transaction manager found for rollback");
            }
            
            targetTM.rollback(status);
            log.debug("Rolled back transaction for data source: {}", dataSourceName);
        } finally {
            cleanup(status);
        }
    }

    /**
     * 初始化事务管理器
     * @param dynamicDataSource 动态数据源
     * @return 数据源名称到事务管理器的映射
     */
    private Map<String, DataSourceTransactionManager> initTransactionManagers(
            DynamicRoutingDataSource dynamicDataSource) {
        Map<String, DataSourceTransactionManager> managers = new ConcurrentHashMap<>();
        
        // 为每个数据源创建对应的事务管理器
        dynamicDataSource.getResolvedDataSources().forEach((name, dataSource) -> {
            try {
                // 包装为XA数据源
                XADataSource xaDataSource = xaWrapper.wrapDataSource(dataSource);
                
                // 创建Atomikos数据源
                AtomikosDataSourceBean atomikosDataSource = new AtomikosDataSourceBean();
                atomikosDataSource.setXaDataSource(xaDataSource);
                atomikosDataSource.setUniqueResourceName(name.toString());
                
                // 创建JTA事务管理器
                managers.put(name.toString(), 
                    new JtaTransactionManager(atomikosDataSource));
                    
                log.info("Initialized transaction manager for data source: {}", name);
            } catch (Exception e) {
                log.error("Failed to initialize transaction manager for: " + name, e);
            }
        });
        
        return managers;
    }

    /**
     * 清理事务资源
     * @param status 事务状态
     */
    private void cleanup(TransactionStatus status) {
        if (status instanceof DefaultTransactionStatus) {
            synchronizationManager.clearSynchronization();
        }
    }
}
 